package com.kedacom.ctsp.authz.security.provider;

import com.google.common.collect.Maps;
import com.kedacom.ctsp.authz.AuthenticationService;
import com.kedacom.ctsp.authz.entity.Authentication;
import com.kedacom.ctsp.authz.token.TokenInformation;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.core.session.SessionDestroyedEvent;
import org.springframework.security.core.session.SessionInformation;
import org.springframework.security.core.session.SessionRegistryImpl;
import org.springframework.util.Assert;

import javax.servlet.http.HttpServletRequest;
import java.util.*;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArraySet;

import static java.util.stream.Collectors.toList;

/**
 * Session In Memory
 *
 * @author xuwei
 * @create 2017-12-04 14:04
 **/
public class MemorySessionRegistryImpl implements AuthzSessionRegistry {

    // ~ Instance fields
    // ================================================================================================

    protected final Log logger = LogFactory.getLog(SessionRegistryImpl.class);

    /**
     * <principal:Object,SessionIdSet>
     */
    private final ConcurrentMap<Object, Set<String>> principals = Maps.newConcurrentMap();
    /**
     * <sessionId:Object,SessionInformation>
     */
    private final Map<String, AuthzSessionInformation> sessionIds = Maps.newConcurrentMap();

    // ~ Methods
    // ========================================================================================================
    @Autowired
    private AuthenticationService authenticationService;

    @Autowired
    private TokenManager tokenManager;

    @Override
    public List<Object> getAllPrincipals() {
        return new ArrayList<>(principals.keySet());
    }

    @Override
    public List<SessionInformation> getAllSessions(Object principal,
                                                   boolean includeExpiredSessions) {
        final Set<String> sessionsUsedByPrincipal = principals.get(principal);

        if (sessionsUsedByPrincipal == null) {
            return Collections.emptyList();
        }

        List<SessionInformation> list = new ArrayList<>(sessionsUsedByPrincipal.size());

        for (String sessionId : sessionsUsedByPrincipal) {
            AuthzSessionInformation sessionInformation = getSessionInformation(sessionId);
            if (sessionInformation == null) {
                continue;
            }
            if (includeExpiredSessions || !sessionInformation.isExpired()) {
                list.add(sessionInformation);
            }
        }

        return list;
    }

    @Override
    public AuthzSessionInformation getSessionInformation(String sessionId) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");

        return sessionIds.get(sessionId);
    }

    @Override
    public void onApplicationEvent(SessionDestroyedEvent event) {
        String sessionId = event.getId();
        removeSessionInformation(sessionId);
    }

    @Override
    public void refreshLastRequest(String sessionId) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");

        SessionInformation info = getSessionInformation(sessionId);

        if (info != null) {
            info.refreshLastRequest();
        }
    }

    @Override
    public void registerNewSession(String sessionId, Object principal) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");
        Assert.notNull(principal, "Principal required as per interface contract");

        if (logger.isDebugEnabled()) {
            logger.debug("Registering session " + sessionId + ", for principal "
                    + principal);
        }

        if (getSessionInformation(sessionId) != null) {
            removeSessionInformation(sessionId);
        }

        sessionIds.put(sessionId,
                new AuthzSessionInformation(principal, sessionId, new Date()));

        Set<String> sessionsUsedByPrincipal = principals.get(principal);

        if (sessionsUsedByPrincipal == null) {
            sessionsUsedByPrincipal = new CopyOnWriteArraySet<String>();
            Set<String> prevSessionsUsedByPrincipal = principals.putIfAbsent(principal,
                    sessionsUsedByPrincipal);
            if (prevSessionsUsedByPrincipal != null) {
                sessionsUsedByPrincipal = prevSessionsUsedByPrincipal;
            }
        }

        sessionsUsedByPrincipal.add(sessionId);

        if (logger.isTraceEnabled()) {
            logger.trace("Sessions used by '" + principal + "' : "
                    + sessionsUsedByPrincipal);
        }
    }

    @Override
    public void removeSessionInformation(String sessionId) {
        Assert.hasText(sessionId, "SessionId required as per interface contract");

        SessionInformation info = getSessionInformation(sessionId);

        if (info == null) {
            return;
        }

        if (logger.isTraceEnabled()) {
            logger.debug("Removing session " + sessionId
                    + " from set of registered sessions");
        }

        sessionIds.remove(sessionId);

        Set<String> sessionsUsedByPrincipal = principals.get(info.getPrincipal());

        if (sessionsUsedByPrincipal == null) {
            return;
        }

        if (logger.isDebugEnabled()) {
            logger.debug("Removing session " + sessionId
                    + " from principal's set of registered sessions");
        }

        sessionsUsedByPrincipal.remove(sessionId);

        if (sessionsUsedByPrincipal.isEmpty()) {
            // No need to keep object in principals Map anymore
            if (logger.isDebugEnabled()) {
                logger.debug("Removing principal " + info.getPrincipal()
                        + " from registry");
            }
            principals.remove(info.getPrincipal());
        }

        if (logger.isTraceEnabled()) {
            logger.trace("Sessions used by '" + info.getPrincipal() + "' : "
                    + sessionsUsedByPrincipal);
        }
    }

    @Override
    public TokenInformation getTokenInfo(String token) {
        return getSessionInformation(token);
    }

    @Override
    public List<TokenInformation> getAllTokens(String username) {
        final Set<String> sessionsUsedByPrincipal = principals.get(new AuthUserDetails(authenticationService.loadUserByUsername(username)));
        if (sessionsUsedByPrincipal == null) {
            return Collections.emptyList();
        }
        List<TokenInformation> list = new ArrayList<>(sessionsUsedByPrincipal.size());
        for (String sessionId : sessionsUsedByPrincipal) {
            AuthzSessionInformation sessionInformation = getSessionInformation(sessionId);
            if (sessionInformation == null) {
                continue;
            }
            if (!sessionInformation.isExpired()) {
                list.add(sessionInformation);
            }
        }
        return list;
    }

    @Override
    public boolean userIsLoggedIn(String username) {
        Set<String> sessionsUsedByPrincipal = principals.get(new AuthUserDetails(authenticationService.loadUserByUsername(username)));
        return CollectionUtils.isNotEmpty(sessionsUsedByPrincipal);
    }

    @Override
    public boolean tokenIsLoggedIn(String token) {
        return getTokenInfo(token) != null;
    }

    @Override
    public long totalUser() {
        return principals.size();
    }

    @Override
    public long totalToken() {
        return sessionIds.size();
    }

    @Override
    public List<Authentication> getAllLoggedUser() {
        return getAllPrincipals().stream()
                .map(AuthUserDetails.class::cast)
                .map(AuthUserDetails::getAuthentication)
                .collect(toList());
    }

    @Override
    public void signOutByUserId(String username) {
        Set<String> sessionsUsedByPrincipal = principals.get(new AuthUserDetails(authenticationService.loadUserByUsername(username)));
        if (CollectionUtils.isNotEmpty(sessionsUsedByPrincipal)) {
            sessionsUsedByPrincipal.stream().map(sessionIds::get).forEach(AuthzSessionInformation::expireNow);
        }

    }

    @Override
    public void signOutByToken(String token) {
        AuthzSessionInformation authzSessionInformation = sessionIds.get(token);
        if (authzSessionInformation != null) {
            authzSessionInformation.expireNow();
        }
    }

    @Override
    public void signIn(HttpServletRequest request, String token, String username) {
        signIn(authenticationService, request, token, username);
    }

    @Override
    public void touch(String token) {
        refreshLastRequest(token);
    }
}